import os
import datetime
import pymysql
from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, Text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from dotenv import load_dotenv

# 全局变量，用于标记数据库连接状态
DB_CONNECTED = False

# 加载环境变量
load_dotenv()

# 数据库配置
DB_HOST = os.getenv("DB_HOST", "localhost")
DB_PORT = int(os.getenv("DB_PORT", "3306"))
DB_USER = os.getenv("DB_USER", "root")
DB_PASSWORD = os.getenv("DB_PASSWORD", "123456")
DB_NAME = os.getenv("DB_NAME", "medical_tubes")

# 创建连接
DATABASE_URL = f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

# 创建基类
Base = declarative_base()

# 定义试管数据模型
class TubeRecord(Base):
    __tablename__ = "tube_records"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    image_path = Column(String(255), nullable=False)
    tube_length = Column(Float, nullable=True)
    cap_color = Column(String(50), nullable=True)
    liquid_position = Column(String(50), nullable=True)
    label_text = Column(Text, nullable=True)
    confidence = Column(Float, nullable=True)
    timestamp = Column(DateTime, default=datetime.datetime.now)
    
    def __repr__(self):
        return f"<TubeRecord(id={self.id}, cap_color={self.cap_color})>"

# 初始化数据库
def init_db():
    global DB_CONNECTED
    try:
        # 检查是否存在数据库，如果不存在则创建
        conn = pymysql.connect(
            host=DB_HOST,
            port=DB_PORT,
            user=DB_USER,
            password=DB_PASSWORD
        )
        cursor = conn.cursor()
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {DB_NAME}")
        conn.commit()
        cursor.close()
        conn.close()
        
        # 创建引擎
        engine = create_engine(DATABASE_URL)
        
        # 创建表
        Base.metadata.create_all(engine)
        DB_CONNECTED = True
        print(f"数据库初始化成功！连接到 {DB_NAME}")
        return True
    except Exception as e:
        DB_CONNECTED = False
        print(f"数据库初始化失败: {str(e)}")
        print("系统将在无数据库模式下运行，数据将不会被保存")
        return False

# 创建会话
def get_session():
    if not DB_CONNECTED:
        print("数据库未连接，无法创建会话")
        return None
        
    try:
        engine = create_engine(DATABASE_URL)
        Session = sessionmaker(bind=engine)
        return Session()
    except Exception as e:
        print(f"创建会话失败: {str(e)}")
        return None

# 保存试管数据
def save_tube_data(image_path, tube_length, cap_color, liquid_position, label_text, confidence):
    if not DB_CONNECTED:
        print("数据库未连接，无法保存数据")
        return None
        
    try:
        session = get_session()
        if not session:
            return None
            
        new_record = TubeRecord(
            image_path=image_path,
            tube_length=tube_length,
            cap_color=cap_color,
            liquid_position=liquid_position,
            label_text=label_text,
            confidence=confidence
        )
        session.add(new_record)
        session.commit()
        record_id = new_record.id
        session.close()
        print(f"数据保存成功，记录ID: {record_id}")
        return record_id
    except Exception as e:
        print(f"保存数据失败: {str(e)}")
        return None

# 获取所有记录
def get_all_records():
    if not DB_CONNECTED:
        print("数据库未连接，无法获取记录")
        return []
        
    try:
        session = get_session()
        if not session:
            return []
            
        records = session.query(TubeRecord).all()
        session.close()
        return records
    except Exception as e:
        print(f"获取记录失败: {str(e)}")
        return []

# 根据ID获取记录
def get_record_by_id(record_id):
    if not DB_CONNECTED:
        print("数据库未连接，无法获取记录")
        return None
        
    try:
        session = get_session()
        if not session:
            return None
            
        record = session.query(TubeRecord).filter(TubeRecord.id == record_id).first()
        session.close()
        return record
    except Exception as e:
        print(f"获取记录失败: {str(e)}")
        return None

# 按条件搜索记录
def search_records(cap_color=None, min_length=None, max_length=None):
    if not DB_CONNECTED:
        print("数据库未连接，无法搜索记录")
        return []
        
    try:
        session = get_session()
        if not session:
            return []
            
        query = session.query(TubeRecord)
        
        if cap_color:
            query = query.filter(TubeRecord.cap_color == cap_color)
        if min_length:
            query = query.filter(TubeRecord.tube_length >= min_length)
        if max_length:
            query = query.filter(TubeRecord.tube_length <= max_length)
            
        records = query.all()
        session.close()
        return records
    except Exception as e:
        print(f"搜索记录失败: {str(e)}")
        return [] 